from flask import Flask, request, render_template,  send_from_directory
import os

from utils import ArgumentParser, ConfigLoader, LOG
from model import   OpenAIModel
from translator import PDFTranslator 

app = Flask(__name__, template_folder='templates')
app.config['UPLOAD_FOLDER'] = 'tests/'   #app.config['UPLOAD_FOLDER'] 定义上传文件夹的路径

app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024  # 设置最大上传大小，这里为32MB  
  
 

@app.route('/')
def index():
    return  render_template('upload.html')

@app.route('/download', methods=['GET'])
def download_request():
    '''
    下载文件
    '''
    try:
         
        # 需要知道2个参数, 第1个参数是本地目录的path, 第2个参数是文件名(带扩展名)
        filename =  request.args.get('filename')  
        directory = os.path.join(app.config['UPLOAD_FOLDER'])  # 假设在当前目录 
        return send_from_directory(directory, filename, as_attachment=True)
    except:
        return '<h1>该文件不存在或无法下载</h1>'
    

@app.route('/translate', methods=['POST'])
def handle_request():
    if request.form.get('model_name') :
        model_name =  request.form.get('model_name') 
    else:
        model_name = app.config['model_name'] 
    target_language = request.form.get('target_language') 
    file_format = request.form.get('file_format') 
    f = request.files['file']  # 提取文件 
    if f.filename == '':  
        return "Please select pdf file", 500 
    if f and allowed_file(f.filename): 
        pdf_file_path = os.path.join(app.config['UPLOAD_FOLDER'],  f.filename)
    
        f.save( pdf_file_path)
        model = OpenAIModel(model=model_name, api_key=app.config['api_key'], base_url=app.config['base_url'])
        translator = PDFTranslator(model)
        output_file_path = translator.translate_pdf(pdf_file_path, file_format,target_language)
        return output_file_path.replace(app.config['UPLOAD_FOLDER'],""), 200  
    else:        
        return "Please upload pdf file", 500  


def allowed_file(filename):  
    """  
    检查文件名是否允许上传  
    这里可以根据需要修改以支持不同的文件类型  
    """  
    ALLOWED_EXTENSIONS = { 'pdf'}  
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS  

if __name__ == '__main__':
    '''
    启动脚本
    python api.py --model_type=OpenAIModel --openai_base_url="*****" --openai_api_key="*******" --openai_model="gpt-3.5-turbo"
    '''
    argument_parser = ArgumentParser()
    args = argument_parser.parse_arguments()
    config_loader = ConfigLoader(args.config)
    
    config = config_loader.load_config()
    app.config['api_key'] = args.openai_api_key if args.openai_api_key else config['OpenAIModel']['api_key']
    app.config['base_url'] = args.openai_base_url if args.openai_base_url else config['OpenAIModel']['base_url']
    app.config['model_name'] = args.openai_model if args.openai_model else config['OpenAIModel']['model']
    app.run(host='0.0.0.0', port=8080,debug=True)